#include <regex>
#include <ShlObj_core.h>

#include "log.h"
#include "crc32c.h"
#include "Share.h"
#include "ShareRecvProcess.h"

// utf-8 text data retrieved from host clipboard with '\0' at the end.
#define RECV_UTF8_TEXT_SIZE_MAX (PLAIN_TEXT_SIZE_MAX + 1)
// HTML data retrieved from host clipboard with '\0' at the end.
#define RECV_HTML_SIZE_MAX (HTML_SIZE_MAX + 1)
// Maximum size of utf-8 text and HTML data.
#define RECV_UTF8_TEXT_AND_HTML_SIZE_MAX (RECV_UTF8_TEXT_SIZE_MAX + RECV_HTML_SIZE_MAX)
// Maximum size of filename List.
#define RECV_FILE_LIST_SIZE_MAX (20 * 1024 * 1024)
// Maximum size of receive header data.
#define RECV_HEADER_SIZE_MAX 64

ShareRecvProcess::~ShareRecvProcess()
{
    m_htmlShareName.clear();
    Clear();
}

BOOL ShareRecvProcess::CheckCrc(ShareHeader *header)
{
    DWORD crcCheck = header->crc;
    header->crc = 0;
    DWORD crcCalc = Crc32c(0, reinterpret_cast<PBYTE>(header), header->headerLen);
    header->crc = crcCheck;
    return crcCalc == crcCheck;
}

BOOL ShareRecvProcess::BasicHeaderCheck(ShareHeader *header)
{
    if (header->signature != SHARE_SIGNATURE) {
        LOGE("Received incorrect signature {}", header->signature);
        return FALSE;
    }

    if (header->headerLen < sizeof(ShareHeader) ||
        header->headerLen > RECV_HEADER_SIZE_MAX) {
        LOGE("Received incorrect header length {}", header->headerLen);
        return FALSE;
    }

    m_dataLen = header->dataLen;
    m_format = header->format;

    return TRUE;
}

BOOL ShareRecvProcess::FormatSpecificCheck(ShareHeader *header)
{
    if (m_format == (m_format & CBSH_UNICODE_AND_HTML)) {
        if (header->headerLen != sizeof(ShareHeader) + sizeof(TextHeader)) {
            LOGE("Received incorrect plain text and HTML header length {}", header->headerLen);
            return FALSE;
        }

        if (m_dataLen > RECV_UTF8_TEXT_AND_HTML_SIZE_MAX) {
            LOGW("Received utf-8 text and HTML size {} exceeds the maximum size {}",
                 m_dataLen, RECV_UTF8_TEXT_AND_HTML_SIZE_MAX);
            return FALSE;
        }
    } else if (m_format == CBSH_FMT_BITMAP) {
        if (header->headerLen != sizeof(ShareHeader)) {
            LOGE("Received incorrect BMP header length {}", header->headerLen);
            return FALSE;
        }

        if (m_dataLen > BITMAP_SIZE_MAX) {
            LOGW("Received BMP size {} exceeds the maximum size {}", m_dataLen, BITMAP_SIZE_MAX);
            return FALSE;
        }
    } else if (m_format == CBSH_FMT_URI_LIST || m_format == CBSH_HTML_FILES_VIRT_SHARE) {
        if (header->headerLen == sizeof(ShareHeader)) {
            LOGE("Received incorrect files virtual share name, format {}", m_format);
            return FALSE;
        }

        if (m_format == CBSH_FMT_URI_LIST && m_dataLen > RECV_FILE_LIST_SIZE_MAX) {
            LOGW("Received file size {} exceeds the maximum size {}", m_dataLen, RECV_FILE_LIST_SIZE_MAX);
            return FALSE;
        }

        if (m_format == CBSH_HTML_FILES_VIRT_SHARE && m_dataLen != 0) {
            LOGE("Received incorrect HTML virtual share name dataLen {}", m_dataLen);
            return FALSE;
        }
    } else if (m_format == CBSH_SYNC_DATA_TO_HOST) {
        if (header->headerLen != sizeof(ShareHeader) || m_dataLen != 0) {
            LOGE("Received incorrect sync clipboard data to host message.");
            return FALSE;
        }
    } else {
        LOGE("Unsupported clipboard format {}", m_format);
        return FALSE;
    }

    return TRUE;
}

BOOL ShareRecvProcess::CheckHeader(ShareHeader *header)
{
    if (!BasicHeaderCheck(header)) {
        return FALSE;
    }

    return FormatSpecificCheck(header);
}

BOOL ShareRecvProcess::CheckExtraHeader()
{
    ShareHeader *header = reinterpret_cast<ShareHeader*>(m_data.data());

    if (!CheckCrc(header)) {
        LOGE("Check clipboard header crc error.");
        return FALSE;
    }

    if (m_format == CBSH_FMT_URI_LIST || m_format == CBSH_HTML_FILES_VIRT_SHARE) {
        if (m_data[header->headerLen - 1] != 0) {
            LOGE("Received incorrect Virtual share name data.");
            return FALSE;
        }

        if (!m_htmlShareName.empty()) {
            m_deleteShareName(m_htmlShareName.c_str());
            m_htmlShareName.clear();
        }

        if (m_format == CBSH_HTML_FILES_VIRT_SHARE) {
            m_htmlShareName = reinterpret_cast<char*>(header + 1);
            m_isHTMLShare = TRUE;
        } else {
            m_fileListShareName = reinterpret_cast<char*>(header + 1);
        }
    } else if (m_format == (m_format & CBSH_UNICODE_AND_HTML)) {
        int ret = memcpy_s(&m_textHeader, sizeof(TextHeader),
                           reinterpret_cast<PBYTE>(header + 1), sizeof(TextHeader));
        if (ret != 0) {
            LOGE("Execution of memcpy_s fails with error: {}", ret);
            return FALSE;
        }

        if (m_textHeader.utf8Len > RECV_UTF8_TEXT_SIZE_MAX ||
            m_textHeader.htmlLen > RECV_HTML_SIZE_MAX ||
            m_textHeader.utf8Len + m_textHeader.htmlLen != header->dataLen) {
            LOGW("Received incorrect extra header, plain text size {}, HTML size {}, all data size {}",
                 m_textHeader.utf8Len, m_textHeader.htmlLen, header->dataLen);
            return FALSE;
        }
    }
    return TRUE;
}

void ShareRecvProcess::Clear()
{
    m_format = 0;
    m_dataLen = 0;
    m_textHeader = {0};
    m_fileListShareName.clear();

    m_data.clear();
    std::vector<BYTE>().swap(m_data);
    m_unicodeText.clear();
    std::vector<WCHAR>().swap(m_unicodeText);
    m_htmlData.clear();
    std::vector<BYTE>().swap(m_htmlData);
}

PBYTE ShareRecvProcess::GetUnicodeData()
{
    return (PBYTE)m_unicodeText.data();
}

SIZE_T ShareRecvProcess::GetUnicodeLen()
{
    return m_unicodeText.size() * sizeof(WCHAR);
}

PBYTE ShareRecvProcess::GetHTMLData()
{
    return (PBYTE)m_htmlData.data();
}

SIZE_T ShareRecvProcess::GetHTMLLen()
{
    return m_htmlData.size();
}

PBYTE ShareRecvProcess::Resize(SIZE_T len)
{
    try {
        m_data.resize(len);
    } catch (const std::exception& e) {
        LOGE("Failed to resize vector: {}", e.what());
        return nullptr;
    }
    return m_data.data();
}

PBYTE ShareRecvProcess::GetData()
{
    return m_data.data();
}

SIZE_T ShareRecvProcess::GetDataLen()
{
    return m_dataLen;
}

DWORD ShareRecvProcess::GetFormat()
{
    return m_format;
}

BOOL ShareRecvProcess::PrepareData()
{
    switch (m_format) {
        case CBSH_FMT_BITMAP:
        case CBSH_FMT_URI_LIST:
        case CBSH_SYNC_DATA_TO_HOST:
            return TRUE;
        case CBSH_FMT_UNICODETEXT:
        case CBSH_FMT_HTML:
        case CBSH_UNICODE_AND_HTML:
             return PrepareText();
        default:
            LOGE("Receive format is not supported: {}", m_format);
            return FALSE;
    }
}

BOOL ShareRecvProcess::PrepareText()
{
    if (m_format & CBSH_FMT_UNICODETEXT) {
        int len = MultiByteToWideChar(CP_UTF8, 0, (LPCCH)m_data.data(), (int)m_textHeader.utf8Len, nullptr, 0);
        if (len <= 0) {
            LOGE("Convert unicode text size failed with error: {}", GetLastError());
            return FALSE;
        }

        m_unicodeText.resize(len);
        if (MultiByteToWideChar(CP_UTF8, 0, (LPCCH)m_data.data(), (int)m_textHeader.utf8Len,
                                (LPWSTR)m_unicodeText.data(), len) <= 0) {
            LOGE("Convert unicode text data failed with error: {}", GetLastError());
            return FALSE;
        }
        m_unicodeText.push_back(0);
    }

    if (m_format & CBSH_FMT_HTML) {
        std::string htmlData;
        if (m_isHTMLShare) {
            std::regex htmlUriRegex(HTML_URI_REGEX);
            std::string replaceHtmlUri = std::string(VIRTIOFSD_TEMP_PATH) + "\\" + m_htmlShareName + "\\";

            htmlData = std::regex_replace(reinterpret_cast<const char*>(m_data.data() + m_textHeader.utf8Len),
                                          htmlUriRegex, "$1" + replaceHtmlUri + "$4");
            m_isHTMLShare = FALSE;
        } else {
            htmlData = reinterpret_cast<const char*>(m_data.data() + m_textHeader.utf8Len);
        }

        /*
        * @StartHtml     - Absolute offset of <html>
        * @EndHtml       - Size of the whole resulting text (excluding ending zero char)
        * @StartFragment - Absolute position after <!--StartFragment-->
        * @EndFragment   - Absolute position of <!--EndFragment-->
        *
        * Note! The offset are zero padded to max width so we don't have any variations due to those.
        * Note! All values includes CRLFs inserted into text.
        *
        * Calculations:
        *   Header length = Format sample length - 2 ('%s')
        *   EndHtml       = Header length + fragment length
        *   StartHtml     = 105(constant)
        *   StartFragment = 141(constant)
        *   EndFragment   = Header length + fragment length - 38 (ending length)
        */
        static const char MSHTML_FORMAT_SAMPLE[] =
            "Version:1.0\r\n"                               /* 13: */
            "StartHTML:0000000105\r\n"                      /* 35: */
            "EndHTML:%00000010u\r\n"                        /* 55: END HTML = Header length + fragment length */
            "StartFragment:0000000141\r\n"                  /* 81: */
            "EndFragment:%00000010u\r\n"                    /* 105: */
            "<html>\r\n<body>\r\n<!--StartFragment-->"      /* 141: */
            "%s"                                            /* 141+2: */
            "<!--EndFragment-->\r\n</body>\r\n</html>\r\n"; /* 179+2: */

        /* Calculate parameters of the MSHTML header */
        const size_t fragmentLen = htmlData.length();
        const size_t headerLen = sizeof(MSHTML_FORMAT_SAMPLE) - 3; /* %s + '\0' */
        const size_t offEndHtml = headerLen + fragmentLen;
        const size_t offEndFragment = headerLen + fragmentLen - 38; /* 179-141 = 38 */

        m_htmlData.resize(offEndHtml + 1);
        int ret = sprintf_s((LPSTR)m_htmlData.data(), offEndHtml + 1,
                            MSHTML_FORMAT_SAMPLE, offEndHtml, offEndFragment, htmlData.c_str());
        if (ret < 0) {
            LOGE("Execution of sprintf_s fails with error: {}", ret);
            return FALSE;
        }
    }
    return TRUE;
}
